/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include "test_common.hpp"
#include <iostream>
#include <vector>
#include <random>
#include <complex>
#include <cmath>
#include <algorithm>
#include <iomanip> // 添加用于复数输出的格式化控制

// ===================== 修复比较函数 =====================
// 解决编译警告和重载问题
template<typename T>
bool near(T a, T b, double eps) { 
    return std::fabs(a - b) < eps; 
}

template<typename T>
bool near(std::complex<T> a, std::complex<T> b, double eps) {
    return std::fabs(a.real() - b.real()) < static_cast<T>(eps) &&
           std::fabs(a.imag() - b.imag()) < static_cast<T>(eps);
}

// ===================== 矩阵生成函数 =====================
template<typename T>
void gen_lower_csc(int n,
                   std::vector<int>& cp,
                   std::vector<int>& ri,
                   std::vector<T>& val) 
{
    std::mt19937 rng(17);
    std::uniform_real_distribution<double> ur(0.2, 1.0);

    cp.resize(n + 1);
    ri.resize(n * (n + 1) / 2);
    val.resize(ri.size());

    int idx = 0;
    for (int j = 0; j < n; ++j) {
        cp[j] = idx;
        for (int r = j; r < n; ++r) {
            ri[idx] = r;
            val[idx] = (r == j) ? static_cast<T>(ur(rng) + 1.0)  // 对角远离0
                                : static_cast<T>(ur(rng));
            ++idx;
        }
    }
    cp[n] = idx;
}

// ===================== CSC矩阵-矩阵乘法 =====================
template<typename T>
void csc_mm(int n, int k,
            const int* cp, const int* ri, const T* v,
            const T* X, int ldx,
            T* Y, int ldy) 
{
    // 初始化结果矩阵
    std::fill(Y, Y + n * k, T{});
    
    for (int colA = 0; colA < n; ++colA) {
        for (int p = cp[colA]; p < cp[colA + 1]; ++p) {
            int row = ri[p];
            const T a = v[p];
            for (int j = 0; j < k; ++j) {
                Y[row + j * ldy] += a * X[colA + j * ldx];
            }
        }
    }
}

// ===================== 改进的输出格式化 =====================
template<typename T>
void print_complex(const std::complex<T>& c) {
    std::cout << "(" << c.real() << ", " << c.imag() << ")";
}

template<typename T>
void print_value(const T& v) {
    if constexpr (std::is_same_v<T, std::complex<float>> || 
                  std::is_same_v<T, std::complex<double>>) {
        print_complex(v);
    } else {
        std::cout << v;
    }
}

// ===================== 测试用例模板 =====================
template<typename T, typename SolveFn>
bool run_case(const char* tag, SolveFn solve, double eps, int n = 16, int nrhs = 2) 
{
    /* 1. 生成测试矩阵 */
    std::vector<int> cp, ri; 
    std::vector<T> val;
    gen_lower_csc(n, cp, ri, val);

    /* 2. 创建随机解矩阵 Xtrue */
    std::mt19937 rng(42);
    std::uniform_real_distribution<double> ur(-1, 1);
    std::vector<T> Xtrue(n * nrhs);
    for (auto& e : Xtrue) e = static_cast<T>(ur(rng));

    /* 3. 计算 B = A · Xtrue */
    std::vector<T> B(n * nrhs);
    csc_mm(n, nrhs, cp.data(), ri.data(), val.data(),
           Xtrue.data(), n, B.data(), n);

    /* 4. 调用三角求解器求解 */
    std::vector<T> X(n * nrhs);
    auto st = solve(n, nrhs,
                    cp.data(), ri.data(), val.data(),
                    X.data(), n,
                    B.data(), n,
                    OPENSPBLAS_ZERO_BASE, OPENSPBLAS_LOWER);

    if (st != OPENSPBLAS_STATUS_SUCCESS) {
        std::cout << tag << " : solver returned error " << st << "\n";
        return false;
    }

    /* 5. 验证结果 */
    bool success = true;
    for (int i = 0; i < n * nrhs; ++i) {
        if (!near(X[i], Xtrue[i], eps)) {
            std::cout << tag << " : mismatch at index " << i << "\n  expected: ";
            print_value(Xtrue[i]);
            std::cout << "\n  actual: ";
            print_value(X[i]);
            std::cout << "\n";
            success = false;
        }
    }
    
    if (success) {
        std::cout << tag << " : PASS\n";
    }
    return success;
}

// ===================== 主函数 =====================
TEST(SparseBLASlevel2, spsm) {
    int fail = 0;
    constexpr int n = 16;
    constexpr int nrhs = 2;
    constexpr float EPS_F = 1e-4f;
    constexpr double EPS_D = 1e-10;

    std::cout << "===== Testing Sparse Triangular Solve for Multiple RHS =====" << std::endl;
    std::cout << "Matrix size: " << n << "x" << n << ", RHS: " << nrhs << std::endl;

    /* float 测试 */
    fail += !run_case<float>(
        "[ssptrsm] (float)",
        [](int n, int k, const int* cp, const int* ri, const float* v,
           float* X, int ldx, const float* B, int ldb,
           OpenSpB_base_type b, OpenSpB_uplo_type u) {
            return OpenSpB_ssptrsm(n, k, cp, ri, v, X, ldx, B, ldb, b, u); 
        },
        EPS_F, n, nrhs);

    /* double 测试 */
    fail += !run_case<double>(
        "[dsptrsm] (double)",
        [](int n, int k, const int* cp, const int* ri, const double* v,
           double* X, int ldx, const double* B, int ldb,
           OpenSpB_base_type b, OpenSpB_uplo_type u) {
            return OpenSpB_dsptrsm(n, k, cp, ri, v, X, ldx, B, ldb, b, u); 
        },
        EPS_D, n, nrhs);

    /* complex<float> 测试 */
    fail += !run_case<std::complex<float>>(
        "[csptrsm] (cfloat)",
        [](int n, int k, const int* cp, const int* ri, const void* v,
           void* X, int ldx, const void* B, int ldb,
           OpenSpB_base_type b, OpenSpB_uplo_type u) {
            return OpenSpB_csptrsm(n, k, cp, ri, v, X, ldx, B, ldb, b, u); 
        },
        EPS_F, n, nrhs);

    /* complex<double> 测试 */
    fail += !run_case<std::complex<double>>(
        "[zsptrsm] (cdouble)",
        [](int n, int k, const int* cp, const int* ri, const void* v,
           void* X, int ldx, const void* B, int ldb,
           OpenSpB_base_type b, OpenSpB_uplo_type u) {
            return OpenSpB_zsptrsm(n, k, cp, ri, v, X, ldx, B, ldb, b, u); 
        },
        EPS_D, n, nrhs);

    /* 测试总结 */
    if (fail == 0) {
        std::cout << "\n===== All sptrsm tests PASSED! =====" << std::endl;
    } else {
        std::cout << "\n===== " << fail << " test(s) FAILED! =====" << std::endl;
    }
}